import torch

# 检查CUDA是否可用
if torch.cuda.is_available():
    print(f"CUDA is available! Using device: {torch.cuda.get_device_name(0)}")
else:
    print("CUDA is not available. Using CPU.")

# 在后续代码中使用 device 变量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")